from core.persona import PersonaManager
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from core.models import PersonaConfig, BasePersonaConfig, ExtendedPersonaConfig
from langchain_core.messages import SystemMessage
from typing import List, Dict, Union

class PromptEngineer:
    def __init__(self, persona_manager: PersonaManager):
        self.persona_manager = persona_manager
    
    def _build_base_persona_prompt(self, persona: BasePersonaConfig) -> str:
        """构建基础人格提示"""
        prompt = f"""
        你正在扮演一位AI心理伙伴，以下是你的角色设定：
        # 角色设定
        ===== AI心理伙伴人设档案 =====
        ■ 角色名称：{persona.name}
        ■ 专业身份：{persona.existing_form}
        ■ 性格特质：{persona.character}
        ■ 沟通风格：{persona.talking_style}
        ■ 人物形象: {persona.portrait}
        ■ 专业背景：{persona.knowledge}
        ■ 交互原则：
        - 标准问候：{persona.greeting}
        - 回退策略：{persona.fallback}
        =================================
        
        """
        return prompt.strip()
    
    def _build_extended_persona_prompt(self, persona: ExtendedPersonaConfig) -> str:
        """构建扩展人格提示"""
        base_prompt = self._build_base_persona_prompt(persona)
        
        extended_prompt = f"""
        {base_prompt}

        # 核心定位
        {self._format_nested_dict(persona.core_positioning)}
        
        # 行为特征
        {self._format_nested_dict(persona.character_behavior)}
        
        # 背景知识
        {self._format_nested_dict(persona.background_knowledge)}
        
        # 功能交互
        {self._format_nested_dict(persona.function_interaction)}
        
        # 扩展定制
        {self._format_nested_dict(persona.expansion_customization)}
        """
        return extended_prompt.strip()
    
    def _format_dict(self, data: Dict[str, str]) -> str:
        """格式化字典为字符串"""
        if not data:
            return ""
        return "\n".join([f"- {k}: {v}" for k, v in data.items()])
    
    def _format_nested_dict(self, data: Dict[str, Union[str, Dict]]) -> str:
        """格式化嵌套字典为字符串"""
        if not data:
            return ""
        
        result = []
        for key, value in data.items():
            if isinstance(value, dict):

                sub_items = self._format_nested_dict(value)
                # sub_items = "\n  ".join([f"{k}: {v}" for k, v in value.items()])
                result.append(f"- {key}:\n  {sub_items}")
            elif isinstance(value, list):
                sub_items = "\n  ".join([f"- {item}" for item in value])
                result.append(f"- {key}:\n  {sub_items}")
            else:
                result.append(f"- {key}: {value}")
        
        return "\n".join(result)
    
    def build_system_message(self, long_term_memory_str=None) -> SystemMessage:
        """根据人格类型构建系统提示"""
        persona = self.persona_manager.get_current_persona()
        
        if persona.is_based:
            content = self._build_base_persona_prompt(persona)
        else:
            content = self._build_extended_persona_prompt(persona)
        
        # 添加通用指导原则
        content += "\n\n# 通用原则\n- 始终以开放式问题结束回应\n- 使用隐喻解释复杂概念\n- 重点在情绪觉察而非解决方案\n- **你的回答对象是用户，所以请不要加一些用括号特别标注的解释, 比如回复的最后一句“（运用心理正常化技术+具体干预方案，保持1提问2建议的节奏，关联用户成功经验）”，这个括号内容是不需要的,不需要回复**\n"

        if long_term_memory_str is not None:
            content += f"\n用户心理背景简介：{long_term_memory_str}\n"

        content += "\n**如果用户输入中有图片表情，可以将图片表情作为参考**\n"
        
        return SystemMessage(content=content)
    
    def build_full_prompt(self, rag_context: List[str] = None) -> ChatPromptTemplate:
        """构建完整对话提示模板"""
        messages = [
            ("system", "{system_message}"),
            MessagesPlaceholder(variable_name="history"),
            # ("human", "{input}")
        ]
        
        if rag_context:
            context_str = "\n".join([f"- {ctx}" for ctx in rag_context])
            messages.insert(1, ("system", f"## 相关背景知识\n{context_str}"))
        
        return ChatPromptTemplate.from_messages(messages)
    
    def build_greeting(self) -> str:
        """构建人格特定的问候语"""
        persona = self.persona_manager.get_current_persona()
        return persona.greeting
    
    def build_fallback_response(self) -> str:
        """构建备用回复"""
        persona = self.persona_manager.get_current_persona()
        return persona.fallback